{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import torch\n",
    "import pytorch3d\n",
    "from pytorch3d.io import IO\n",
    "from urdfpy import URDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "urdf_path = '/private/home/ssax/hab_spot_arm/urdf/hab_spot_arm.urdf'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for link in robot.joints:\n",
    "    print(link._name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urdfpy\n",
    "import trimesh\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "# Load the URDF file\n",
    "robot = urdfpy.URDF.load(urdf_path)\n",
    "\n",
    "# Create an empty Trimesh object to store the mesh\n",
    "mesh = trimesh.Trimesh()\n",
    "\n",
    "# # Iterate through all the visual elements in the URDF\n",
    "# for link in robot.links:\n",
    "#     for visual in link.visuals:\n",
    "#         # Check if the visual element has a mesh\n",
    "#         print(visual.geometry.__dict__)\n",
    "#         if visual.geometry.type == 'mesh':\n",
    "#             # Load the mesh from the file specified in the URDF\n",
    "#             mesh_file = visual.geometry.filename\n",
    "#             mesh_data = trimesh.load_mesh(mesh_file)\n",
    "#             # Transform the mesh to the visual element's pose if needed\n",
    "#             # For example, you might need to apply the link's transform\n",
    "#             # to position the mesh correctly\n",
    "#             mesh_data.apply_transform(visual.origin)\n",
    "\n",
    "#             # Add the mesh to the combined mesh\n",
    "#             mesh += mesh_data\n",
    "\n",
    "# Iterate through all the visual elements in the URDF\n",
    "# Compute visual trimesh forward kinematics\n",
    "cfg = {}\n",
    "for leg in ['fl', 'fr', 'hl', 'hr']:\n",
    "    cfg[f'{leg}.hy'] = 1*np.pi/6\n",
    "    cfg[f'{leg}.kn'] = -2*np.pi/6\n",
    "     \n",
    "visuals = robot.visual_trimesh_fk(cfg={\n",
    "    'arm0.hr0': np.pi,\n",
    "    'arm0.sh0': - np.pi,\n",
    "    'arm0.sh1': - np.pi * 1 / 8,\n",
    "    'arm0.el0': np.pi * 5 / 8,\n",
    "    'arm0.wr0': np.pi * 2 / 8 + np.pi / 12,\n",
    "    # 'arm0.wr1': np.pi,\n",
    "    'arm0.f1x': -np.pi * 1 / 4 - np.pi / 12,\n",
    "    **cfg \n",
    "})\n",
    "\n",
    "# Initialize an empty list to hold the transformed meshes\n",
    "meshes = []\n",
    "\n",
    "# Iterate over each visual in the visuals dictionary\n",
    "for mesh, transform in visuals.items():\n",
    "    mesh.apply_transform(transform)\n",
    "    meshes.append(mesh)\n",
    "\n",
    "# Combine the meshes\n",
    "print(len(meshes))\n",
    "combined_mesh = trimesh.util.concatenate(meshes)\n",
    "\n",
    "# Export the combined mesh to an OBJ file\n",
    "output_obj_path = \"output_robot.obj\"\n",
    "combined_mesh.export(output_obj_path, file_type=\"obj\", include_texture=True)\n",
    "\n",
    "print(f\"OBJ file saved to {output_obj_path}\")\n",
    "\n",
    "# from PIL import Image\n",
    "\n",
    "# # Define a factor by which to lighten the image\n",
    "# lighten_factor = 100  # this value can be adjusted to your needs\n",
    "\n",
    "# # Open the image file\n",
    "# with Image.open('material_0.png') as img:\n",
    "#     img = img.convert('RGBA')  # Ensure the image has an alpha channel\n",
    "#     data = img.getdata()\n",
    "    \n",
    "#     # Create a new data list to hold the modified pixels\n",
    "#     new_data = []\n",
    "#     for item in data:\n",
    "#         # Change all white (also shades of whites)\n",
    "#         # pixels to be fully transparent\n",
    "#         r, g, b, a = item\n",
    "#         new_data.append((r + lighten_factor, g + lighten_factor, b + lighten_factor, 255))\n",
    "        \n",
    "#     # Update image data\n",
    "#     img.putdata(new_data)\n",
    "    \n",
    "#     # Save new image\n",
    "#     img.save('material_0.png')\n",
    "\n",
    "combined_mesh.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch3d.renderer import (\n",
    "    TexturesAtlas,\n",
    "    TexturesVertex,\n",
    "    TexturesUV,\n",
    ")\n",
    "from pytorch3d.io import IO, load_obj, load_ply\n",
    "\n",
    "out_path = 'output_robot.obj'\n",
    "out_path = Path(out_path)\n",
    "os.chdir(out_path.parent)\n",
    "with open(Path(out_path), 'r') as f:\n",
    "    mesh = IO().load_mesh(out_path, device='cpu')\n",
    "    print(f\"Loaded mesh with {len(mesh.faces_packed())} faces\")\n",
    "\n",
    "def convert_to_textureVertex(textures_uv: TexturesUV, meshes:pytorch3d.structures.Meshes) -> TexturesVertex:\n",
    "    verts_colors_packed = torch.zeros_like(meshes.verts_packed())\n",
    "    verts_colors_packed[meshes.faces_packed()] = textures_uv.faces_verts_textures_packed()  # (*)\n",
    "    return TexturesVertex(pytorch3d.structures.packed_to_list(verts_colors_packed, meshes.num_verts_per_mesh()))\n",
    "\n",
    "if isinstance(mesh.textures, TexturesUV):\n",
    "    mesh.textures = convert_to_textureVertex(mesh.textures, mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# os.chdir('assets/spot_mesh/')\n",
    "\n",
    "# mesh = IO().load_mesh('assets/spot_mesh/output_robot.obj')\n",
    "from pytorch3d.io import load_objs_as_meshes, load_obj\n",
    "from pytorch3d.vis.plotly_vis import plot_scene\n",
    "plot_scene(\n",
    "    {'scene': {\n",
    "        'mesh': mesh\n",
    "        }\n",
    "     }\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "home-robot118",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
